#pragma once
#pragma once

#include "clipper/clipper.hpp"
// locality-aware NMS
namespace lanms {

	namespace cl = ClipperLib;

	struct Polygon {
		cl::Path poly;
		float score;
	};

	float paths_area(const ClipperLib::Paths& ps) {
		float area = 0;
		for (auto&& p : ps)
			area += cl::Area(p);
		return area;
	}

	float poly_iou(const Polygon& a, const Polygon& b) {
		cl::Clipper clpr;
		clpr.AddPath(a.poly, cl::ptSubject, true);
		clpr.AddPath(b.poly, cl::ptClip, true);

		cl::Paths inter, uni;
		clpr.Execute(cl::ctIntersection, inter, cl::pftEvenOdd);
		clpr.Execute(cl::ctUnion, uni, cl::pftEvenOdd);

		auto inter_area = paths_area(inter),
			uni_area = paths_area(uni);
		return std::abs(inter_area) / std::max(std::abs(uni_area), 1.0f);
	}

	bool should_merge(const Polygon& a, const Polygon& b, float iou_threshold) {
		return poly_iou(a, b) > iou_threshold;
	}

	/**
	 * Incrementally merge polygons
	 */
	class PolyMerger {
	public:
		PolyMerger() : score(0), nr_polys(0) {
			memset(data, 0, sizeof(data));
		}

		/**
		 * Add a new polygon to be merged.
		 */
		void add(const Polygon& p_given) {
			Polygon p;
			if (nr_polys > 0) {
				// vertices of two polygons to merge may not in the same order;
				// we match their vertices by choosing the ordering that
				// minimizes the total squared distance.
				// see function normalize_poly for details.
				p = normalize_poly(get(), p_given);
			}
			else {
				p = p_given;
			}
			assert(p.poly.size() == 4);
			auto& poly = p.poly;
			auto s = p.score;
			data[0] += poly[0].X * s;
			data[1] += poly[0].Y * s;

			data[2] += poly[1].X * s;
			data[3] += poly[1].Y * s;

			data[4] += poly[2].X * s;
			data[5] += poly[2].Y * s;

			data[6] += poly[3].X * s;
			data[7] += poly[3].Y * s;

			score += p.score;

			nr_polys += 1;
		}

		inline std::int64_t sqr(std::int64_t x) { return x * x; }

		Polygon normalize_poly(
			const Polygon& ref,
			const Polygon& p) {

			std::int64_t min_d = std::numeric_limits<std::int64_t>::max();
			size_t best_start = 0, best_order = 0;

			for (size_t start = 0; start < 4; start++) {
				size_t j = start;
				std::int64_t d = (
					sqr(ref.poly[(j + 0) % 4].X - p.poly[(j + 0) % 4].X)
					+ sqr(ref.poly[(j + 0) % 4].Y - p.poly[(j + 0) % 4].Y)
					+ sqr(ref.poly[(j + 1) % 4].X - p.poly[(j + 1) % 4].X)
					+ sqr(ref.poly[(j + 1) % 4].Y - p.poly[(j + 1) % 4].Y)
					+ sqr(ref.poly[(j + 2) % 4].X - p.poly[(j + 2) % 4].X)
					+ sqr(ref.poly[(j + 2) % 4].Y - p.poly[(j + 2) % 4].Y)
					+ sqr(ref.poly[(j + 3) % 4].X - p.poly[(j + 3) % 4].X)
					+ sqr(ref.poly[(j + 3) % 4].Y - p.poly[(j + 3) % 4].Y)
					);
				if (d < min_d) {
					min_d = d;
					best_start = start;
					best_order = 0;
				}

				d = (
					sqr(ref.poly[(j + 0) % 4].X - p.poly[(j + 3) % 4].X)
					+ sqr(ref.poly[(j + 0) % 4].Y - p.poly[(j + 3) % 4].Y)
					+ sqr(ref.poly[(j + 1) % 4].X - p.poly[(j + 2) % 4].X)
					+ sqr(ref.poly[(j + 1) % 4].Y - p.poly[(j + 2) % 4].Y)
					+ sqr(ref.poly[(j + 2) % 4].X - p.poly[(j + 1) % 4].X)
					+ sqr(ref.poly[(j + 2) % 4].Y - p.poly[(j + 1) % 4].Y)
					+ sqr(ref.poly[(j + 3) % 4].X - p.poly[(j + 0) % 4].X)
					+ sqr(ref.poly[(j + 3) % 4].Y - p.poly[(j + 0) % 4].Y)
					);
				if (d < min_d) {
					min_d = d;
					best_start = start;
					best_order = 1;
				}
			}

			Polygon r;
			r.poly.resize(4);
			auto j = best_start;
			if (best_order == 0) {
				for (size_t i = 0; i < 4; i++)
					r.poly[i] = p.poly[(j + i) % 4];
			}
			else {
				for (size_t i = 0; i < 4; i++)
					r.poly[i] = p.poly[(j + 4 - i - 1) % 4];
			}
			r.score = p.score;
			return r;
		}

		Polygon get() const {
			Polygon p;

			auto& poly = p.poly;
			poly.resize(4);
			auto score_inv = 1.0f / std::max(1e-8f, score);
			poly[0].X = data[0] * score_inv;
			poly[0].Y = data[1] * score_inv;
			poly[1].X = data[2] * score_inv;
			poly[1].Y = data[3] * score_inv;
			poly[2].X = data[4] * score_inv;
			poly[2].Y = data[5] * score_inv;
			poly[3].X = data[6] * score_inv;
			poly[3].Y = data[7] * score_inv;

			assert(score > 0);
			p.score = score;

			return p;
		}

	private:
		std::int64_t data[8];
		float score;
		std::int32_t nr_polys;
	};


	/**
	 * The standard NMS algorithm.
	 */
	std::vector<Polygon> standard_nms(std::vector<Polygon>& polys, float iou_threshold) {
		size_t n = polys.size();
		if (n == 0)
			return {};
		std::vector<size_t> indices(n);
		std::iota(std::begin(indices), std::end(indices), 0);
		std::sort(std::begin(indices), std::end(indices), [&](size_t i, size_t j) { return polys[i].score > polys[j].score; });

		std::vector<size_t> keep;
		while (indices.size()) {
			size_t p = 0, cur = indices[0];
			keep.emplace_back(cur);
			for (size_t i = 1; i < indices.size(); i++) {
				if (!should_merge(polys[cur], polys[indices[i]], iou_threshold)) {
					indices[p++] = indices[i];
				}
			}
			indices.resize(p);
		}

		std::vector<Polygon> ret;
		for (auto&& i : keep) {
			ret.emplace_back(polys[i]);
		}
		return ret;
	}

	std::vector<Polygon>
		merge_quadrangle_n9(std::vector<std::vector<float>> data, float iou_threshold, float precision) {
		using cInt = cl::cInt;

		// first pass
		std::vector<Polygon> polys;
		for (auto &p : data) {
			Polygon poly{
				{
					{cInt(p[0] * precision), cInt(p[1] * precision)},
					{cInt(p[2] * precision), cInt(p[3] * precision)},
					{cInt(p[4] * precision), cInt(p[5] * precision)},
					{cInt(p[6] * precision), cInt(p[7] * precision)},
				},
				p[8],
			};

			if (polys.size()) {
				// merge with the last one
				auto& bpoly = polys.back();
				if (should_merge(poly, bpoly, iou_threshold)) {
					PolyMerger merger;
					merger.add(bpoly);
					merger.add(poly);
					bpoly = merger.get();
				}
				else {
					polys.emplace_back(poly);
				}
			}
			else {
				polys.emplace_back(poly);
			}
		}
		return standard_nms(polys, iou_threshold);
	}

	std::vector<std::vector<float>> polys2floats(const std::vector<lanms::Polygon>& polys, float precision) {
		std::vector<std::vector<float>> ret;
		for (size_t i = 0; i < polys.size(); i++) {
			auto& p = polys[i];
			auto& poly = p.poly;
			ret.emplace_back(std::vector<float>{
				float(poly[0].X / precision), float(poly[0].Y / precision),
					float(poly[1].X / precision), float(poly[1].Y / precision),
					float(poly[2].X / precision), float(poly[2].Y / precision),
					float(poly[3].X / precision), float(poly[3].Y / precision)
			});
		}

		return ret;
	}
}
